import scipy.io as scio
import mcubes
import numpy as np
import os
from PyCmpltrtok.common import sep, get_dir_name_ext, rand_name_on_now

np.random.seed(88888)

sep('Prepare for saving')
xdir, xbase, xext = get_dir_name_ext(os.path.abspath(__file__))
save_dir = os.path.join(xdir, '_save', xbase)
os.makedirs(save_dir, exist_ok=True)

sep('Load bfm')
bfm = scio.loadmat('/var/asuspei/large_data/face/bfm/__workspace/PublicMM1/01_MorphableModel.mat')

sep('Calc random face')
n_dim = 4
n_seg = 1
factor = -5.0

coef_type = 'ones'
coef = np.ones([n_dim, n_seg], dtype=np.float32) * factor
# coef_type = '0001'
# coef = np.array([0, 0, 0, 1], dtype=np.float32).reshape(n_dim, n_seg) * factor
mu = bfm['shapeMU']
pc = bfm['shapePC']
ev = bfm['shapeEV']

base = mu @ np.ones([1, n_seg], dtype=np.float32)
coef_mat = ev[:n_dim, :] @ np.ones([1, n_seg], dtype=np.float32) * coef
variance = pc[:, :n_dim] @ coef_mat
face_result = base + variance

sep('bfm2dae')
vertices = face_result
triangles = bfm['tl']

vertices = vertices.reshape(-1, 3)
# vertices = np.transpose(vertices, [1, 0])
vertices = vertices.copy()
print('vertices', vertices.shape)

triangles = triangles.copy()
triangles -= 1
print('triangles', triangles.shape)

sep('Extract dae')
dae_path = os.path.join(save_dir, f'specified_bfm2dae_coef_{coef_type}_factor{factor}_ndim{n_dim}_.dt{rand_name_on_now()}.dae')
print(f'Extracting to {dae_path}')
mcubes.export_mesh(vertices, triangles, dae_path, "BFM2dae")
print(f'Extracted.')

sep('All over')
